from sklearn.tree import DecisionTreeClassifier
from common import get_train_label
from scipy import sparse
import pickle
import os


# 训练决策树模型
def trainDecisionTreeModel(name, path, header, sep=','):
    if not os.path.exists('../model/DecisionTree/%s_DT_MODEL.pkl' % name):
        # 使用默认的配置对决策树分类器进行初始化。
        dtCls = DecisionTreeClassifier()
        train, label = get_train_label(path, header, sep=sep)
        train_fit = sparse.load_npz("../model/matrix/%s_TF_IDF_TRAIN_FIT_MATRIX.npz" % name)
        # 使用朴素决策树分类器，对Tfidf后的训练样本进行参数学习。
        dtCls.fit(train_fit, label)
        # 将训练好的决策树模型保存到文件
        pickle.dump(dtCls, open("../model/DecisionTree/%s_DT_MODEL.pkl" % name, "wb"))
    else:
        print(name + "数据 DecisionTree 模型已进行训练，若需重新训练，删除文件即可")


if __name__ == '__main__':
    # 以微博数据训练决策树模型
    trainDecisionTreeModel(name='weibo', path='../data/weibo_senti_100k.csv', header=['label', 'review'])
    print('微博数据决策树模型训练完毕')
    # 以外卖数据训练决策树模型
    trainDecisionTreeModel(name='waimai', path='../data/waimai_10k.csv', header=['label', 'review'])
    print('外卖数据决策树模型训练完毕')
    # 以电商数据训练决策树模型
    trainDecisionTreeModel(name='shopping', path='../data/online_shopping_10_cats.csv',
                           header=['cat', 'label', 'review'])
    print('电商数据决策树模型训练完毕')
    # 以训练数据训练决策树模型
    trainDecisionTreeModel(name='train', path='../data/train.txt', header=['label', 'review'], sep='\t')
    print('训练数据决策树模型训练完毕')
